# 引入库
import os
import time
import signal

import pyaudio
from math import sqrt
import numpy as np
import audioop

import struct

from multiprocessing import Process, Event, Queue, freeze_support

#############################################################
from modbus_tk import modbus_rtu
import serial

OS_ENV = os.getenv("OS_ENV", "local")

LIGHT_HEIGHT = 100
CHUNKS_PER_SECOND = 60
SAMPLE_RATE = int(str(os.getenv("SAMPLE_RATE", 44100)))  # 16000
CHUNK = int(str(os.getenv("CHUNK", int((1 / CHUNKS_PER_SECOND) * SAMPLE_RATE))))
INDEX_OF_CHOSEN_INPUT_DEVICE = int(str(os.getenv("INDEX_OF_CHOSEN_INPUT_DEVICE", 1)))
INDEX_OF_CHOSEN_OUTPUT_DEVICE = int(str(os.getenv("INDEX_OF_CHOSEN_OUTPUT_DEVICE", 1)))
MODE_PLAYER = os.getenv("MODE_PLAYER", "playback")

BEAT_THRESHOLD = int(str(os.getenv("BEAT_THRESHOLD", 60)))

if (OS_ENV == "container"):
    TCP_SERVER_ADDRESS = os.getenv("STEL_TCP_SERVER_ADDR", "192.168.1.200")
    TCP_SERVER_PORT = int(str(os.getenv("STEL_TCP_SERVER_PORT", 8000)))
    config = {"ip": "172.17.0.1", "port": 30555, "protocol": "grpc", "insecure": True}
elif (OS_ENV == "local"):
    TCP_SERVER_ADDRESS = os.getenv("STEL_TCP_SERVER_ADDR", "192.168.211.148")
    TCP_SERVER_PORT = int(str(os.getenv("STEL_TCP_SERVER_PORT", 8080)))
    config = {"ip": "127.0.0.1", "port": 55555, "protocol": "grpc", "insecure": True}


def term(sig_num, addtion):
    print('term current pid is %s, group id is %s' % (os.getpid(), os.getpgrp()))
    os.killpg(os.getpgid(os.getpid()), signal.SIGKILL)


signal.signal(signal.SIGTERM, term)


def is_beat(peak, threshold=80):
    return peak > threshold


def getPeak(wav_data, type):
    peak = 0
    if (type == "fft"):
        waveform = np.frombuffer(wav_data, dtype=np.int16)
        fft_complex = np.fft.fft(waveform, n=CHUNK)
        # these operations can be factored out.
        # sqrt does not need to be computed for every distance
        max_val = sqrt(max(v.real * v.real + v.imag * v.imag for v in fft_complex))
        # print("max_val=", max_val)
        # factor out the scale multiply.
        scale_value = LIGHT_HEIGHT / max_val
        # print("scale_value=", scale_value)
        # for index and value in the frequencies
        for i, v in enumerate(fft_complex):
            # calculate the distance using the ecludian distance measure.
            dist = sqrt(v.real * v.real + v.imag * v.imag)
            peak = dist * scale_value
    elif (type == "rms"):
        peak = audioop.rms(wav_data, 2)

    return int(peak)


class Devs:
    def __init__(self):
        self.__color = {"R": 0, "G": 0, "B": 0}
        self.__wiper = 0
        self.__highBeam = 0
        self.__leftTurnLight = 0
        self.__rightTurnLight = 0
        self.__brakeLight = 0

        self.__colorAudio = {"R": 0, "G": 0, "B": 0}
        self.__wiperAudio = 0
        self.__highBeamAudio = 0
        self.__leftTurnLightAudio = 0
        self.__rightTurnLightAudio = 0
        self.__brakeLightAudio = 0

        self.stopNumber = 0
        self.startNumber = 0
        self.playStatus = False

        self.devs_init()

    def devs_init(self):
        self.serial = serial.Serial(port='/dev/ttyUSB0', baudrate=115200)
        print("devs init")

    def devs_send(self, cmd: str, protol="rs232"):

        cmdBytes = bytes.fromhex(cmd)

        cmdLlen = len(cmdBytes)
        cmdLen = '{:02x}'.format(cmdLlen)
        cmdAT = f'ff {cmdLen} {cmd}'
        # print("cmdAT=", cmdAT)

        cmdBytes = bytes.fromhex(cmdAT)
        if (protol == "rs232"):
            self.serial.write(cmdBytes)
        else:
            print("other protol")

    def checkPlayStatus(self, peak: int):
        if (peak < 10):
            self.stopNumber = self.stopNumber + 1
            if (self.stopNumber > 50):
                self.startNumber = 0
                self.stopNumber = 0
                self.playStatus = False
        else:
            self.startNumber = self.startNumber + 1
            if (self.startNumber > 50):
                self.startNumber = 0
                self.stopNumber = 0
                self.playStatus = True

    # color: r - [0, 254], g - [0, 254],b - [0, 254],
    def __setColor(self, colorValue: dict, protol="rs232"):
        if (self.__colorAudio == colorValue):
            return

        self.__colorAudio = colorValue

        light = colorValue

        light["S"] = 1  # cmd
        s = '{:02x}'.format(int(light["S"]))
        r = '{:02x}'.format(int(light["R"]))
        g = '{:02x}'.format(int(light["G"]))
        b = '{:02x}'.format(int(light["B"]))

        checksum = int(light["S"]) + int(light["R"]) + int(light["G"]) + int(light["B"])
        if (checksum > 255):
            checksum = '{:02x}'.format(int(str(checksum)[-2:]))
        else:
            checksum = '{:02x}'.format(checksum)

        cmd = f'{s} {r} {g} {b} {checksum}'

        self.devs_send(cmd)

    # action: 0 - stop, 1 - low, 2 -med, 3 - high
    def __setWiper(self, action: int, protol="rs232"):
        if (self.__wiperAudio == action):
            return

        self.__wiperAudio = action

        wiperCmd = 2  # cmd
        wiperSpeed = action

        c = '{:02x}'.format(wiperCmd)
        s = '{:02x}'.format(wiperSpeed)
        checksum = wiperCmd + wiperSpeed

        if (checksum > 255):
            checksum = '{:02x}'.format(int(str(checksum)[-2:]))
        else:
            checksum = '{:02x}'.format(checksum)

        cmd = f'{c} {s} {checksum}'

        self.devs_send(cmd)

    # action: 0 - close, 1 - open
    def __setHighBeam(self, action: int, protol="rs232"):
        if (self.__highBeamAudio == action):
            return
        self.__highBeamAudio = action

        cmd = 3  # cmd

        c = '{:02x}'.format(cmd)
        s = '{:02x}'.format(action)
        checksum = cmd + action

        if (checksum > 255):
            checksum = '{:02x}'.format(int(str(checksum)[-2:]))
        else:
            checksum = '{:02x}'.format(checksum)

        cmd = f'{c} {s} {checksum}'

        self.devs_send(cmd)

    # action: 0 - close, 1 - open
    def __setLeftTurnLight(self, action: int, protol="rs232"):
        if (self.__leftTurnLightAudio == action):
            return
        self.__leftTurnLightAudio = action

        cmd = 5  # cmd

        c = '{:02x}'.format(cmd)
        s = '{:02x}'.format(action)
        checksum = cmd + action

        if (checksum > 255):
            checksum = '{:02x}'.format(int(str(checksum)[-2:]))
        else:
            checksum = '{:02x}'.format(checksum)

        cmd = f'{c} {s} {checksum}'

        self.devs_send(cmd)

    def __setRightTurnLight(self, action: int, protol="rs232"):
        if (self.__rightTurnLightAudio == action):
            return

        self.__rightTurnLightAudio = action

        cmd = 6  # cmd

        c = '{:02x}'.format(cmd)
        s = '{:02x}'.format(action)
        checksum = cmd + action

        if (checksum > 255):
            checksum = '{:02x}'.format(int(str(checksum)[-2:]))
        else:
            checksum = '{:02x}'.format(checksum)

        cmd = f'{c} {s} {checksum}'

        self.devs_send(cmd)

        # action: 0 - close, 1 - open

    def __setBrakeLight(self, action: int, protol="rs232"):
        if (self.__brakeLightAudio == action):
            return
        self.__brakeLightAudio = action

        cmd = 7  # cmd

        c = '{:02x}'.format(cmd)
        s = '{:02x}'.format(action)
        checksum = cmd + action

        if (checksum > 255):
            checksum = '{:02x}'.format(int(str(checksum)[-2:]))
        else:
            checksum = '{:02x}'.format(checksum)

        cmd = f'{c} {s} {checksum}'

        self.devs_send(cmd)

    def devsControlByAudio(self, audioBeatQue: Queue, colorQue: Queue, wiperQue: Queue, highBeamQue: Queue,
                           leftTurnLightQue: Queue, rightTurnLightQue: Queue, brakeLightQue: Queue):
        print("Devices Control By Audio...")

        while True:
            # light
            if not colorQue.empty():
                color = colorQue.get()
                self.__color = color
                self.__setColor(color)

            # wiper
            if not wiperQue.empty():
                wiper = wiperQue.get()
                self.__wiper = wiper
                self.__setWiper(wiper)

            # highBeam
            if not highBeamQue.empty():
                highBeam = highBeamQue.get()
                self.__highBeam = highBeam
                self.__setHighBeam(highBeam)

            if not leftTurnLightQue.empty():
                leftTurnLight = leftTurnLightQue.get()
                self.__leftTurnLight = leftTurnLight
                self.__setLeftTurnLight(leftTurnLight)

            if not rightTurnLightQue.empty():
                rightTurnLight = rightTurnLightQue.get()
                self.__rightTurnLight = rightTurnLight
                self.__setRightTurnLight(rightTurnLight)

            if not brakeLightQue.empty():
                brakeLight = brakeLightQue.get()
                self.__brakeLight = brakeLight
                self.__setBrakeLight(brakeLight)

            # audioLight
            if not audioBeatQue.empty():
                audioData = audioBeatQue.get()

                peak = getPeak(audioData, "fft")

                self.checkPlayStatus(peak)
                # print("peak playStatus", peak, self.playStatus, self.__color, self.__highBeam, self.__leftTurnLight, self.__rightTurnLight, self.__brakeLight)

                if (self.playStatus == False):
                    self.__setColor(self.__color)  # keep
                    self.__setWiper(self.__wiper)
                    self.__setHighBeam(self.__highBeam)
                    self.__setLeftTurnLight(self.__leftTurnLight)
                    self.__setRightTurnLight(self.__rightTurnLight)
                    self.__setBrakeLight(self.__brakeLight)
                else:
                    # print("peak=",peak)
                    if (is_beat(peak, threshold=BEAT_THRESHOLD)):
                        self.__setColor(self.__color)  # keep
                        self.__setWiper(3)
                        self.__setHighBeam(1)
                        self.__setLeftTurnLight(1)
                        self.__setRightTurnLight(1)
                        self.__setBrakeLight(1)
                    else:
                        color = {"R": 0, "G": 0, "B": 0}
                        self.__setColor(color)
                        self.__setWiper(0)
                        self.__setHighBeam(0)
                        self.__setLeftTurnLight(0)
                        self.__setRightTurnLight(0)
                        self.__setBrakeLight(0)

                    time.sleep(0.01)

    def start(self, beatQue: Queue, colorQue: Queue, wiperQue: Queue, highBeamQue: Queue, leftTurnLightQue: Queue,
              rightTurnLightQue: Queue, brakeLightQue: Queue):
        print("devs start")

        devs_process = Process(target=self.devsControlByAudio, args=(
        beatQue, colorQue, wiperQue, highBeamQue, leftTurnLightQue, rightTurnLightQue, brakeLightQue))
        devs_process.start()


class Audio:
    # init function
    def __init__(self):
        self.playback_que = Queue(2048)
        self.event = Event()

        self.audio_init()

    def getAudioIndex(self):
        audio_index = 0
        p = pyaudio.PyAudio()
        info = p.get_host_api_info_by_index(0)
        numdevices = info.get('deviceCount')
        for i in range(0, numdevices):
            if (p.get_device_info_by_host_api_device_index(0, i).get('maxInputChannels')) > 0:
                print("Input Device id ", i, " - ", p.get_device_info_by_host_api_device_index(0, i))
                device_info = p.get_device_info_by_host_api_device_index(0, i)
                print("device_info[name] =", device_info["name"])
                if 'USB Audio Device' in device_info["name"]:
                    audio_index = i
        print("audio_index=", audio_index)
        return audio_index

    def audio_init(self):
        self.pformat = pyaudio.paInt16
        self.pchannels = 1
        self.prate = SAMPLE_RATE
        self.input_device = self.getAudioIndex()
        self.output_device = self.getAudioIndex()
        self.chunk = CHUNK

    def playback(self, que):
        print("playback...")
        audio = pyaudio.PyAudio()
        streamOut = audio.open(format=self.pformat,
                               channels=self.pchannels,
                               rate=self.prate,
                               output=True,
                               output_device_index=self.output_device,
                               frames_per_buffer=self.chunk)

        while True:
            if not que.empty():
                data = que.get()
                streamOut.write(data)

    def listen(self, playbackQue: Queue, beatQue: Queue, audioQue: Queue):
        print("listen...")
        print("self.pformat=", self.pformat)
        print("self.prate=", self.prate)
        print("self.chunk=", self.chunk)
        print("self.pchannels=", self.pchannels)

        audio = pyaudio.PyAudio()
        streamIn = audio.open(format=self.pformat,
                              channels=self.pchannels,
                              rate=self.prate,
                              input=True,
                              input_device_index=self.input_device,
                              frames_per_buffer=self.chunk)

        while True:
            audio_data = streamIn.read(self.chunk, exception_on_overflow=False)
            peak = getPeak(audio_data, "fft")
            if (peak > 0):
                playbackQue.put(audio_data)
                if audioQue.qsize() < 1024:
                    audioQue.put(peak)
                beatQue.put(audio_data)
                # print("qsize=",audioQue.qsize(), beatQue.qsize(), playbackQue.qsize())

    def start(self, audioQue, beatQue):
        print("start")

        # 监听
        listen_process = Process(target=self.listen, args=(self.playback_que, beatQue, audioQue))
        listen_process.start()

        # 播放
        play_process = Process(target=self.playback, args=(self.playback_que,))
        play_process.start()


class LightPlayer():
    def __init__(self):
        freeze_support()
        self.audio_que = Queue(2048)
        self.beat_que = Queue(2048)
        self.color_que = Queue(2048)
        self.wiper_que = Queue(2048)
        self.highBeam_que = Queue(2048)
        self.brakeLight_que = Queue(2048)
        self.leftTurnLight_que = Queue(2048)
        self.rightTurnLight_que = Queue(2048)

        self.setLightColor({"R": 60, "G": 60, "B": 60})

    def getAudioStream(self) -> bytes:
        return struct.pack('i', self.audio_que.get())

    def setLightColor(self, color: dict):
        self.color_que.put(color)

    def setWiper(self, speed: int):
        self.wiper_que.put(speed)

    def setHighBeam(self, action: int):
        self.highBeam_que.put(action)

    def setLeftTurnLight(self, action: int):
        self.leftTurnLight_que.put(action)

    def setRightTurnLight(self, action: int):
        self.rightTurnLight_que.put(action)

    def setBrakeLight(self, action: int):
        self.brakeLight_que.put(action)

    def start(self):
        audio = Audio()
        devs = Devs()
        audio.start(self.audio_que, self.beat_que)
        devs.start(self.beat_que, self.color_que, self.wiper_que, self.highBeam_que, self.leftTurnLight_que,
                   self.rightTurnLight_que, self.brakeLight_que)


lightPlayer = LightPlayer()
